#include "websocket_server.h"
#include "logger.h"
#include <iostream>
#include <sstream>
#include <algorithm>
#include <cstring>
#include <random>
#include <cerrno>
#include <openssl/sha.h>
#include <openssl/bio.h>
#include <openssl/evp.h>
#include <openssl/buffer.h>

#ifdef _WIN32
#pragma comment(lib, "ws2_32.lib")
#else
#include <sys/select.h>
#include <fcntl.h>
#endif

namespace yquote {

WebSocketServer::WebSocketServer(int port) 
    : port_(port), server_socket_(-1), running_(false), heartbeat_interval_(30) {
    initializeNetwork();
}

WebSocketServer::~WebSocketServer() {
    stop();
    cleanupNetwork();
}

bool WebSocketServer::start() {
    if (running_.load()) {
        return true; // 已经在运行
    }
    
    // 创建服务器套接字
    server_socket_ = socket(AF_INET, SOCK_STREAM, 0);
    if (server_socket_ < 0) {
        LOG_ERROR("创建套接字失败");
        return false;
    }
    
    // 设置套接字选项
    int opt = 1;
#ifdef _WIN32
    setsockopt(server_socket_, SOL_SOCKET, SO_REUSEADDR, (char*)&opt, sizeof(opt));
#else
    setsockopt(server_socket_, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
#endif
    
    // 绑定地址
    sockaddr_in server_addr{};
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(port_);
    
    if (bind(server_socket_, (sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
        LOG_ERROR_FMT("绑定端口 %d 失败", port_);
#ifdef _WIN32
        closesocket(server_socket_);
#else
        close(server_socket_);
#endif
        return false;
    }
    
    // 开始监听
    if (listen(server_socket_, 10) < 0) {
        LOG_ERROR("监听失败");
#ifdef _WIN32
        closesocket(server_socket_);
#else
        close(server_socket_);
#endif
        return false;
    }
    
    running_.store(true);
    accept_thread_ = std::thread(&WebSocketServer::acceptLoop, this);
    heartbeat_thread_ = std::thread(&WebSocketServer::heartbeatLoop, this);
    
    LOG_INFO_FMT("WebSocket服务器已启动，监听端口: %d", port_);
    return true;
}

void WebSocketServer::stop() {
    if (!running_.load()) {
        return;
    }
    
    LOG_INFO("正在停止WebSocket服务器...");
    running_.store(false);
    
    // 关闭服务器套接字
    if (server_socket_ >= 0) {
#ifdef _WIN32
        closesocket(server_socket_);
#else
        close(server_socket_);
#endif
        server_socket_ = -1;
    }
    
    // 收集所有需要等待的线程，避免在持有锁时等待
    std::vector<std::thread*> client_threads;
    
    // 先关闭所有客户端连接，让阻塞的recv调用返回
    {
        std::lock_guard<std::mutex> lock(clients_mutex_);
        for (auto& pair : clients_) {
            auto& client = pair.second;
            client->state.store(ClientConnectionState::CLOSED);
#ifdef _WIN32
            closesocket(client->socket);
#else
            close(client->socket);
#endif
            // 收集线程指针，稍后在锁外等待
            if (client->thread.joinable()) {
                client_threads.push_back(&client->thread);
            }
        }
    }
    
    // 等待accept和heartbeat线程结束
    if (accept_thread_.joinable()) {
        accept_thread_.join();
    }
    if (heartbeat_thread_.joinable()) {
        heartbeat_thread_.join();
    }
    
    // 在锁外等待所有客户端线程结束，避免死锁
    for (auto* thread_ptr : client_threads) {
        if (thread_ptr->joinable()) {
            thread_ptr->join();
        }
    }
    
    // 最后清理客户端映射
    {
        std::lock_guard<std::mutex> lock(clients_mutex_);
        clients_.clear();
    }
    
    LOG_INFO("WebSocket服务器已停止");
}

bool WebSocketServer::isRunning() const {
    return running_.load();
}

bool WebSocketServer::sendMessage(const std::string& client_id, const std::string& message) {
    std::lock_guard<std::mutex> lock(clients_mutex_);
    
    auto it = clients_.find(client_id);
    if (it != clients_.end() && it->second->state.load() == ClientConnectionState::CONNECTED) {
        std::lock_guard<std::mutex> send_lock(it->second->send_mutex);
        return sendWebSocketFrame(it->second->socket, message);
    }
    
    return false;
}

bool WebSocketServer::sendBinaryMessage(const std::string& client_id, const std::vector<uint8_t>& data) {
    std::lock_guard<std::mutex> lock(clients_mutex_);
    
    auto it = clients_.find(client_id);
    if (it != clients_.end() && it->second->state.load() == ClientConnectionState::CONNECTED) {
        std::lock_guard<std::mutex> send_lock(it->second->send_mutex);
        return sendWebSocketFrame(it->second->socket, WebSocketFrameType::BINARY, data);
    }
    
    return false;
}

void WebSocketServer::broadcast(const std::string& message) {
    std::lock_guard<std::mutex> lock(clients_mutex_);
    
    for (const auto& pair : clients_) {
        if (pair.second->state.load() == ClientConnectionState::CONNECTED) {
            std::lock_guard<std::mutex> send_lock(pair.second->send_mutex);
            sendWebSocketFrame(pair.second->socket, message);
        }
    }
}

void WebSocketServer::broadcastBinary(const std::vector<uint8_t>& data) {
    std::lock_guard<std::mutex> lock(clients_mutex_);
    
    for (const auto& pair : clients_) {
        if (pair.second->state.load() == ClientConnectionState::CONNECTED) {
            std::lock_guard<std::mutex> send_lock(pair.second->send_mutex);
            sendWebSocketFrame(pair.second->socket, WebSocketFrameType::BINARY, data);
        }
    }
}

void WebSocketServer::sendToSubscribers(const std::vector<std::string>& client_ids, const std::string& message) {
    std::lock_guard<std::mutex> lock(clients_mutex_);
    
    for (const auto& client_id : client_ids) {
        auto it = clients_.find(client_id);
        if (it != clients_.end() && it->second->state.load() == ClientConnectionState::CONNECTED) {
            std::lock_guard<std::mutex> send_lock(it->second->send_mutex);
            sendWebSocketFrame(it->second->socket, message);
        }
    }
}

bool WebSocketServer::sendPing(const std::string& client_id, const std::vector<uint8_t>& payload) {
    std::lock_guard<std::mutex> lock(clients_mutex_);
    
    auto it = clients_.find(client_id);
    if (it != clients_.end() && it->second->state.load() == ClientConnectionState::CONNECTED) {
        std::lock_guard<std::mutex> send_lock(it->second->send_mutex);
        it->second->last_ping_time = std::chrono::steady_clock::now();
        return sendWebSocketFrame(it->second->socket, WebSocketFrameType::PING, payload);
    }
    
    return false;
}

bool WebSocketServer::closeClient(const std::string& client_id, uint16_t code, const std::string& reason) {
    std::lock_guard<std::mutex> lock(clients_mutex_);
    
    auto it = clients_.find(client_id);
    if (it != clients_.end() && it->second->state.load() == ClientConnectionState::CONNECTED) {
        it->second->state.store(ClientConnectionState::CLOSING);
        std::lock_guard<std::mutex> send_lock(it->second->send_mutex);
        return sendCloseFrame(it->second->socket, code, reason);
    }
    
    return false;
}

void WebSocketServer::setMessageCallback(MessageCallback callback) {
    message_callback_ = std::move(callback);
}

void WebSocketServer::setClientConnectCallback(ClientConnectCallback callback) {
    connect_callback_ = std::move(callback);
}

void WebSocketServer::setClientDisconnectCallback(ClientDisconnectCallback callback) {
    disconnect_callback_ = std::move(callback);
}

void WebSocketServer::setErrorCallback(ErrorCallback callback) {
    error_callback_ = std::move(callback);
}

size_t WebSocketServer::getClientCount() const {
    std::lock_guard<std::mutex> lock(clients_mutex_);
    return clients_.size();
}

std::vector<std::string> WebSocketServer::getClientIds() const {
    std::lock_guard<std::mutex> lock(clients_mutex_);
    
    std::vector<std::string> ids;
    ids.reserve(clients_.size());
    
    for (const auto& pair : clients_) {
        ids.push_back(pair.first);
    }
    
    return ids;
}

ClientConnectionState WebSocketServer::getClientState(const std::string& client_id) const {
    std::lock_guard<std::mutex> lock(clients_mutex_);
    
    auto it = clients_.find(client_id);
    if (it != clients_.end()) {
        return it->second->state.load();
    }
    
    return ClientConnectionState::CLOSED;
}

void WebSocketServer::setHeartbeatInterval(int seconds) {
    heartbeat_interval_ = seconds;
}

void WebSocketServer::acceptLoop() {
    while (running_.load()) {
        // 使用select实现带超时的accept
        fd_set read_fds;
        FD_ZERO(&read_fds);
        FD_SET(server_socket_, &read_fds);
        
        // 设置超时时间为1秒
        struct timeval timeout;
        timeout.tv_sec = 1;
        timeout.tv_usec = 0;
        
        int select_result = select(server_socket_ + 1, &read_fds, nullptr, nullptr, &timeout);
        
        if (select_result < 0) {
            if (running_.load()) {
                LOG_ERROR_FMT("select失败: %s", strerror(errno));
            }
            continue;
        }
        
        if (select_result == 0) {
            // 超时，继续循环检查running_状态
            continue;
        }
        
        if (!running_.load()) {
            break;
        }
        
        // 有新连接可以接受
        if (FD_ISSET(server_socket_, &read_fds)) {
            sockaddr_in client_addr{};
            socklen_t client_len = sizeof(client_addr);
            
            int client_socket = accept(server_socket_, (sockaddr*)&client_addr, &client_len);
            if (client_socket < 0) {
                if (running_.load()) {
                    LOG_ERROR_FMT("接受连接失败: %s", strerror(errno));
                }
                continue;
            }
            
            // 执行WebSocket握手
            if (!performHandshake(client_socket)) {
                LOG_ERROR("WebSocket握手失败");
#ifdef _WIN32
                closesocket(client_socket);
#else
                close(client_socket);
#endif
                continue;
            }
            
            // 创建客户端连接
            std::string client_id = generateClientId();
            auto client = std::make_unique<ClientConnection>(client_socket, client_id);
            client->state.store(ClientConnectionState::CONNECTED);
            
            // 先将客户端添加到映射中
            {
                std::lock_guard<std::mutex> lock(clients_mutex_);
                clients_[client_id] = std::move(client);
            }
            
            // 启动客户端处理线程
            {
                std::lock_guard<std::mutex> lock(clients_mutex_);
                clients_[client_id]->thread = std::thread(&WebSocketServer::handleClient, this, client_id);
            }
            
            LOG_INFO_FMT("新客户端连接: %s", client_id.c_str());
        }
    }
}

void WebSocketServer::heartbeatLoop() {
    const int check_interval_ms = 1000; // 每秒检查一次
    int elapsed_seconds = 0;
    
    while (running_.load()) {
        // 短时间睡眠，便于快速响应关闭信号
        std::this_thread::sleep_for(std::chrono::milliseconds(check_interval_ms));
        
        if (!running_.load()) break;
        
        elapsed_seconds++;
        
        // 只有达到心跳间隔时才执行心跳检查
        if (elapsed_seconds < heartbeat_interval_) {
            continue;
        }
        
        // 重置计数器
        elapsed_seconds = 0;
        
        auto now = std::chrono::steady_clock::now();
        std::vector<std::string> clients_to_remove;
        
        {
            std::lock_guard<std::mutex> lock(clients_mutex_);
            for (const auto& pair : clients_) {
                const auto& client = pair.second;
                if (client->state.load() != ClientConnectionState::CONNECTED) {
                    continue;
                }
                
                // 检查是否需要发送心跳
                auto ping_elapsed = std::chrono::duration_cast<std::chrono::seconds>(
                    now - client->last_ping_time).count();
                
                if (ping_elapsed >= heartbeat_interval_) {
                    // 发送Ping帧
                    std::lock_guard<std::mutex> send_lock(client->send_mutex);
                    if (sendWebSocketFrame(client->socket, WebSocketFrameType::PING, {})) {
                        client->last_ping_time = now;
                        LOG_DEBUG_FMT("向客户端 %s 发送心跳", pair.first.c_str());
                    } else {
                        clients_to_remove.push_back(pair.first);
                    }
                }
                
                // 检查心跳超时
                auto pong_elapsed = std::chrono::duration_cast<std::chrono::seconds>(
                    now - client->last_pong_time).count();
                
                if (pong_elapsed >= heartbeat_interval_ * 2) {
                    LOG_WARN_FMT("客户端心跳超时: %s", pair.first.c_str());
                    clients_to_remove.push_back(pair.first);
                }
            }
        }
        
        // 移除超时的客户端
        for (const auto& client_id : clients_to_remove) {
            removeClient(client_id);
        }
    }
}

void WebSocketServer::handleClient(const std::string& client_id) {
    // 通知客户端连接
    if (connect_callback_) {
        connect_callback_(client_id);
    }
    
    // 处理客户端消息
    while (running_.load()) {
        ClientConnection* client_ptr = nullptr;
        int client_socket = -1;
        
        // 快速获取客户端信息，减少锁持有时间
        {
            std::lock_guard<std::mutex> lock(clients_mutex_);
            auto it = clients_.find(client_id);
            if (it != clients_.end() && it->second->state.load() != ClientConnectionState::CLOSED) {
                client_ptr = it->second.get();
                client_socket = client_ptr->socket;
            } else {
                break; // 客户端已断开或不活跃
            }
        }
        
        // 检查socket是否有效
        if (client_socket < 0) {
            break;
        }
        
        WebSocketFrame frame;
        std::string message;
        WebSocketReceiveResult result = receiveWebSocketFrame(client_socket, frame, message);
        
        switch (result) {
            case WebSocketReceiveResult::SUCCESS:
                // 处理业务消息
                if (message_callback_ && !message.empty()) {
                    message_callback_(client_id, message);
                }
                break;
                
            case WebSocketReceiveResult::PING_RECEIVED:
                // Ping帧已经在receiveWebSocketFrame中自动回复了Pong
                LOG_DEBUG_FMT("收到客户端 %s 的Ping帧", client_id.c_str());
                break;
                
            case WebSocketReceiveResult::PONG_RECEIVED:
                // 更新最后收到Pong的时间
                {
                    std::lock_guard<std::mutex> lock(clients_mutex_);
                    auto it = clients_.find(client_id);
                    if (it != clients_.end()) {
                        it->second->last_pong_time = std::chrono::steady_clock::now();
                    }
                }
                LOG_DEBUG_FMT("收到客户端 %s 的Pong帧", client_id.c_str());
                break;
                
            case WebSocketReceiveResult::CLOSE_RECEIVED:
                LOG_INFO_FMT("收到客户端 %s 的关闭帧", client_id.c_str());
                goto cleanup; // 跳出循环进行清理
                
            case WebSocketReceiveResult::CONNECTION_ERROR:
            case WebSocketReceiveResult::PROTOCOL_ERROR:
                reportError(client_id, "连接或协议错误");
                goto cleanup; // 跳出循环进行清理
                
            case WebSocketReceiveResult::UNSUPPORTED_FRAME:
                LOG_DEBUG("收到不支持的帧类型，忽略");
                break;
                
            case WebSocketReceiveResult::NO_DATA:
                // 非阻塞模式下没有数据，继续等待
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
                break;
        }
    }
    
cleanup:
    // 标记客户端为已关闭并关闭socket
    {
        std::lock_guard<std::mutex> lock(clients_mutex_);
        auto it = clients_.find(client_id);
        if (it != clients_.end()) {
            it->second->state.store(ClientConnectionState::CLOSED);
            // 只有在socket还没被关闭时才关闭
            if (it->second->socket >= 0) {
#ifdef _WIN32
                closesocket(it->second->socket);
#else
                close(it->second->socket);
#endif
                it->second->socket = -1; // 标记为已关闭
            }
        }
    }
    
    // 通知客户端断开连接
    if (disconnect_callback_) {
        disconnect_callback_(client_id);
    }
    
    LOG_INFO_FMT("客户端断开连接: %s", client_id.c_str());
}

bool WebSocketServer::performHandshake(int client_socket) {
    // 简化的WebSocket握手实现
    char buffer[4096];
    int bytes_received = recv(client_socket, buffer, sizeof(buffer) - 1, 0);
    if (bytes_received <= 0) {
        return false;
    }
    
    buffer[bytes_received] = '\0';
    std::string request(buffer);
    
    // 检查是否是WebSocket升级请求
    if (request.find("Upgrade: websocket") == std::string::npos) {
        return false;
    }
    
    // 提取Sec-WebSocket-Key
    size_t key_pos = request.find("Sec-WebSocket-Key: ");
    if (key_pos == std::string::npos) {
        return false;
    }
    
    key_pos += 19; // "Sec-WebSocket-Key: " 的长度
    size_t key_end = request.find("\r\n", key_pos);
    if (key_end == std::string::npos) {
        return false;
    }
    
    std::string key = request.substr(key_pos, key_end - key_pos);
    
    // 计算Sec-WebSocket-Accept
    std::string accept_key = calculateWebSocketAccept(key);
    
    // 生成正确的WebSocket握手响应
    std::string response = 
        "HTTP/1.1 101 Switching Protocols\r\n"
        "Upgrade: websocket\r\n"
        "Connection: Upgrade\r\n"
        "Sec-WebSocket-Accept: " + accept_key + "\r\n"
        "\r\n";
    
    return send(client_socket, response.c_str(), response.length(), 0) > 0;
}

std::string WebSocketServer::calculateWebSocketAccept(const std::string& key) {
    // WebSocket协议规定的GUID
    const std::string websocket_guid = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
    
    // 连接key和GUID
    std::string combined = key + websocket_guid;
    
    // 计算SHA-1哈希
    unsigned char hash[SHA_DIGEST_LENGTH];
    SHA1(reinterpret_cast<const unsigned char*>(combined.c_str()), combined.length(), hash);
    
    // Base64编码
    BIO* bio = BIO_new(BIO_s_mem());
    BIO* b64 = BIO_new(BIO_f_base64());
    BIO_set_flags(b64, BIO_FLAGS_BASE64_NO_NL);
    bio = BIO_push(b64, bio);
    
    BIO_write(bio, hash, SHA_DIGEST_LENGTH);
    BIO_flush(bio);
    
    BUF_MEM* buffer_ptr;
    BIO_get_mem_ptr(bio, &buffer_ptr);
    
    std::string result(buffer_ptr->data, buffer_ptr->length);
    BIO_free_all(bio);
    
    return result;
}

bool WebSocketServer::sendWebSocketFrame(int socket, const WebSocketFrame& frame) {
    std::vector<uint8_t> frame_data;
    
    // 第一个字节：FIN + RSV + Opcode
    uint8_t first_byte = 0;
    if (frame.fin) first_byte |= 0x80;
    if (frame.rsv1) first_byte |= 0x40;
    if (frame.rsv2) first_byte |= 0x20;
    if (frame.rsv3) first_byte |= 0x10;
    first_byte |= static_cast<uint8_t>(frame.opcode) & 0x0F;
    frame_data.push_back(first_byte);
    
    // 第二个字节：MASK + Payload length
    uint8_t second_byte = 0;
    if (frame.masked) second_byte |= 0x80;
    
    if (frame.payload_length < 126) {
        second_byte |= static_cast<uint8_t>(frame.payload_length);
        frame_data.push_back(second_byte);
    } else if (frame.payload_length < 65536) {
        second_byte |= 126;
        frame_data.push_back(second_byte);
        frame_data.push_back(static_cast<uint8_t>((frame.payload_length >> 8) & 0xFF));
        frame_data.push_back(static_cast<uint8_t>(frame.payload_length & 0xFF));
    } else {
        second_byte |= 127;
        frame_data.push_back(second_byte);
        for (int i = 7; i >= 0; --i) {
            frame_data.push_back(static_cast<uint8_t>((frame.payload_length >> (i * 8)) & 0xFF));
        }
    }
    
    // 掩码（服务器发送的帧通常不需要掩码）
    if (frame.masked) {
        for (int i = 0; i < 4; ++i) {
            frame_data.push_back(frame.mask[i]);
        }
    }
    
    // 负载数据
    if (frame.masked) {
        for (size_t i = 0; i < frame.payload.size(); ++i) {
            frame_data.push_back(frame.payload[i] ^ frame.mask[i % 4]);
        }
    } else {
        frame_data.insert(frame_data.end(), frame.payload.begin(), frame.payload.end());
    }
    
    return safeSend(socket, frame_data.data(), frame_data.size());
}

bool WebSocketServer::sendWebSocketFrame(int socket, WebSocketFrameType type, const std::vector<uint8_t>& data) {
    WebSocketFrame frame;
    frame.fin = true;
    frame.opcode = type;
    frame.masked = false; // 服务器发送的帧不需要掩码
    frame.payload_length = data.size();
    frame.payload = data;
    
    return sendWebSocketFrame(socket, frame);
}

bool WebSocketServer::sendWebSocketFrame(int socket, const std::string& data) {
    std::vector<uint8_t> payload(data.begin(), data.end());
    return sendWebSocketFrame(socket, WebSocketFrameType::TEXT, payload);
}

WebSocketReceiveResult WebSocketServer::receiveWebSocketFrame(int socket, WebSocketFrame& frame, std::string& message) {
    // 解析帧头
    if (!parseFrameHeader(socket, frame)) {
        return WebSocketReceiveResult::CONNECTION_ERROR;
    }
    
    // 读取负载数据
    if (!readFramePayload(socket, frame)) {
        return WebSocketReceiveResult::CONNECTION_ERROR;
    }
    
    // 处理不同类型的帧
    switch (frame.opcode) {
        case WebSocketFrameType::TEXT:
            message = std::string(frame.payload.begin(), frame.payload.end());
            return WebSocketReceiveResult::SUCCESS;
            
        case WebSocketFrameType::BINARY:
            // 对于二进制数据，可以考虑使用base64编码或其他方式处理
            message = std::string(frame.payload.begin(), frame.payload.end());
            return WebSocketReceiveResult::SUCCESS;
            
        case WebSocketFrameType::PING:
            // 自动回复Pong帧
            sendPongFrame(socket, frame.payload);
            return WebSocketReceiveResult::PING_RECEIVED;
            
        case WebSocketFrameType::PONG:
            return WebSocketReceiveResult::PONG_RECEIVED;
            
        case WebSocketFrameType::CLOSE:
            return WebSocketReceiveResult::CLOSE_RECEIVED;
            
        case WebSocketFrameType::CONTINUATION:
            // 暂不支持分片消息
            return WebSocketReceiveResult::UNSUPPORTED_FRAME;
            
        default:
            return WebSocketReceiveResult::UNSUPPORTED_FRAME;
    }
}

bool WebSocketServer::parseFrameHeader(int socket, WebSocketFrame& frame) {
    uint8_t header[2];
    if (!safeRecv(socket, header, 2)) {
        return false;
    }
    
    // 解析第一个字节
    frame.fin = (header[0] & 0x80) != 0;
    frame.rsv1 = (header[0] & 0x40) != 0;
    frame.rsv2 = (header[0] & 0x20) != 0;
    frame.rsv3 = (header[0] & 0x10) != 0;
    frame.opcode = static_cast<WebSocketFrameType>(header[0] & 0x0F);
    
    // 解析第二个字节
    frame.masked = (header[1] & 0x80) != 0;
    frame.payload_length = header[1] & 0x7F;
    
    // 处理扩展长度
    if (frame.payload_length == 126) {
        uint8_t len_bytes[2];
        if (!safeRecv(socket, len_bytes, 2)) {
            return false;
        }
        frame.payload_length = (len_bytes[0] << 8) | len_bytes[1];
    } else if (frame.payload_length == 127) {
        uint8_t len_bytes[8];
        if (!safeRecv(socket, len_bytes, 8)) {
            return false;
        }
        frame.payload_length = 0;
        for (int i = 0; i < 8; ++i) {
            frame.payload_length = (frame.payload_length << 8) | len_bytes[i];
        }
        
        // 检查长度是否合理（避免过大的帧）
        if (frame.payload_length > 1024 * 1024) { // 1MB限制
            return false;
        }
    }
    
    // 读取掩码
    if (frame.masked) {
        if (!safeRecv(socket, frame.mask, 4)) {
            return false;
        }
    }
    
    return true;
}

bool WebSocketServer::readFramePayload(int socket, WebSocketFrame& frame) {
    if (frame.payload_length == 0) {
        return true;
    }
    
    frame.payload.resize(frame.payload_length);
    if (!safeRecv(socket, frame.payload.data(), frame.payload_length)) {
        return false;
    }
    
    // 解除掩码
    if (frame.masked) {
        for (uint64_t i = 0; i < frame.payload_length; ++i) {
            frame.payload[i] ^= frame.mask[i % 4];
        }
    }
    
    return true;
}

bool WebSocketServer::sendPongFrame(int socket, const std::vector<uint8_t>& payload) {
    return sendWebSocketFrame(socket, WebSocketFrameType::PONG, payload);
}

bool WebSocketServer::sendCloseFrame(int socket, uint16_t code, const std::string& reason) {
    std::vector<uint8_t> payload;
    
    // 添加关闭码
    payload.push_back(static_cast<uint8_t>((code >> 8) & 0xFF));
    payload.push_back(static_cast<uint8_t>(code & 0xFF));
    
    // 添加关闭原因
    payload.insert(payload.end(), reason.begin(), reason.end());
    
    return sendWebSocketFrame(socket, WebSocketFrameType::CLOSE, payload);
}

bool WebSocketServer::safeRecv(int socket, void* buffer, size_t size) {
    size_t total_received = 0;
    char* buf = static_cast<char*>(buffer);
    
    while (total_received < size && running_.load()) {
        // 使用select检查socket是否有数据可读，设置超时
        fd_set read_fds;
        FD_ZERO(&read_fds);
        FD_SET(socket, &read_fds);
        
        struct timeval timeout;
        timeout.tv_sec = 1;  // 1秒超时
        timeout.tv_usec = 0;
        
        int select_result = select(socket + 1, &read_fds, nullptr, nullptr, &timeout);
        
        if (select_result < 0) {
            // select错误
            return false;
        }
        
        if (select_result == 0) {
            // 超时，继续循环检查running_状态
            continue;
        }
        
        if (!running_.load()) {
            // 服务器正在关闭
            return false;
        }
        
        // 有数据可读
        if (FD_ISSET(socket, &read_fds)) {
            ssize_t received = recv(socket, buf + total_received, size - total_received, 0);
            if (received <= 0) {
                return false;
            }
            total_received += received;
        }
    }
    
    return total_received == size;
}

bool WebSocketServer::safeSend(int socket, const void* buffer, size_t size) {
    size_t total_sent = 0;
    const char* buf = static_cast<const char*>(buffer);
    
    while (total_sent < size) {
        ssize_t sent = send(socket, buf + total_sent, size - total_sent, 0);
        if (sent <= 0) {
            return false;
        }
        total_sent += sent;
    }
    
    return true;
}

void WebSocketServer::reportError(const std::string& client_id, const std::string& error) {
    if (error_callback_) {
        error_callback_(client_id, error);
    }
    LOG_ERROR_FMT("客户端 %s 错误: %s", client_id.c_str(), error.c_str());
}

std::string WebSocketServer::generateClientId() {
    static std::random_device rd;
    static std::mt19937 gen(rd());
    static std::uniform_int_distribution<> dis(0, 15);
    
    std::string id = "client_";
    for (int i = 0; i < 8; ++i) {
        id += "0123456789abcdef"[dis(gen)];
    }
    
    return id;
}

bool WebSocketServer::initializeNetwork() {
#ifdef _WIN32
    WSADATA wsaData;
    return WSAStartup(MAKEWORD(2, 2), &wsaData) == 0;
#else
    return true;
#endif
}

void WebSocketServer::cleanupNetwork() {
#ifdef _WIN32
    WSACleanup();
#endif
}

void WebSocketServer::removeClient(const std::string& client_id) {
    std::unique_ptr<ClientConnection> client;
    
    {
        std::lock_guard<std::mutex> lock(clients_mutex_);
        auto it = clients_.find(client_id);
        if (it != clients_.end()) {
            client = std::move(it->second);
            clients_.erase(it);
        }
    }
    
    if (client) {
        client->state.store(ClientConnectionState::CLOSED);
#ifdef _WIN32
        closesocket(client->socket);
#else
        close(client->socket);
#endif
        
        // 只有当不是当前线程时才join
        if (client->thread.joinable() && client->thread.get_id() != std::this_thread::get_id()) {
            client->thread.join();
        } else if (client->thread.joinable()) {
            client->thread.detach(); // 如果是当前线程，则分离
        }
        
        // 通知客户端断开连接
        if (disconnect_callback_) {
            disconnect_callback_(client_id);
        }
        
        LOG_INFO_FMT("客户端断开连接: %s", client_id.c_str());
    }
}

} // namespace yquote 